

from ast import arg
from cProfile import label
import copy

import cv2
import numpy as np

import torch
import torch.nn.functional as F
# import torchvision

from my_py_toolkit.file.file_toolkit import *
from my_py_toolkit.log.logger import get_logger
from PIL import Image
from torch.optim import Adam
from tqdm import tqdm
from lib.args import get_args_sbi
from lib.data import get_dataloader, get_dataloader_hj, get_dataloader_sbi
from lib.model import ClsEff
from lib.utils import *
from lib.evaluate import *

    
    
def main():
    
    args = get_args_sbi()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_ids
    args.gpu_ids = list(range(len(args.gpu_ids.split(','))))
    logger = get_logger(f'{args.log_dir}/{args.log_file}', console_out=False)
    model = ClsEff(2).to(args.device)
    if args.resume:
        logger.info(f'start resume epoch : {args.resume_epoch}' + '-' * 80)
        model.load_state_dict(torch.load(f'{args.model_dir}/model_{args.resume_epoch}.pth'))       
    if len(args.gpu_ids) > 1:
        model = torch.nn.DataParallel(model, args.gpu_ids)
    params = [p for p in model.parameters() if p.requires_grad]
    opt = Adam(params, lr=args.lr)
    dataloader = get_dataloader_sbi(args.train_dir, args.landmarks_path, args.train_faces_path, args.batch_size // 2, 
                                args.img_size, 'train')
    test_dataloader = get_dataloader_sbi(args.test_dir, args.landmarks_path, args.test_faces_path, args.batch_size, 
                                args.img_size, 'test')
    
    
    writer = get_writer(args) 
    if args.mode == 'train':
        train(args, model, opt, dataloader, test_dataloader, writer, logger)
    elif args.mode == 'test':
        evaluate_pr(model, test_dataloader, args, writer, 10000)


def train(args, model, opt, dataloader, test_dataloader, writer, logger):
    steps_global = 0
    for epoch in range(args.epochs):
        model.train()
        train_bar = tqdm(range(len(dataloader)))
        for i, (imgs, labels) in enumerate(dataloader):
            opt.zero_grad()
            imgs, labels = imgs.to(args.device), labels.to(args.device)
            loss = model(imgs, labels)
            loss.backward()
            opt.step()
            train_bar.set_description(f'epoch {epoch} loss {loss.item()}')
            writer.add_scalar('loss', loss.item(), steps_global)
            logger.info(f'epoch {epoch} loss {loss.item()}')
            train_bar.update(1)
            steps_global += 1
            if steps_global % args.steps_for_eval == 0:
                evaluate_pr(model, test_dataloader, args, writer, steps_global, args.steps_eval)
                model.train()
            
        torch.save(model.state_dict(), f'{args.model_dir}/model_{epoch}.pth')
        train_bar.close()



if __name__ == "__main__":
    main()